package com.xu.rag.service.impl;

import cn.hutool.core.io.FastByteArrayOutputStream;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.io.IoUtil;
import cn.hutool.core.util.StrUtil;
import com.xu.rag.common.enums.ModelEnum;
import com.xu.rag.factory.ModelFactory;
import com.xu.rag.pojo.req.ChatReq;
import com.xu.rag.pojo.result.ChunkResult;
import com.xu.rag.pojo.result.EmbeddingResult;
import com.xu.rag.service.LlmService;
import com.xu.rag.service.TestService;
import com.xu.rag.service.VectorService;
import com.xu.rag.utils.LlmUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;

import javax.annotation.Resource;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @Author: xdl
 * @CreateTime: 2025-01-19 21:10
 * @Description:
 */
@Slf4j
@Service
public class TestServiceImpl implements TestService {

    @Resource
    private VectorService vectorService;

    /**
     * 暂且就搞个txt
     * @param file
     * @return
     */
    @Override
    public String embedding(MultipartFile file) {
        try {
            String txtContent = IoUtil.read(file.getInputStream(), StandardCharsets.UTF_8);
            List<ChunkResult> chunkResults = chunk(file.getOriginalFilename(), txtContent);
            LlmService tongyiService = ModelFactory.getModel(ModelEnum.TONG_YI.getCode());
            List<EmbeddingResult> embeddingResults = tongyiService.embedding(chunkResults);
            String collectionName = vectorService.getCollectionName();
            vectorService.storage(collectionName, embeddingResults);
        } catch (IOException e) {
            log.error("embedding error {}", e.getMessage());
        }
        return "SUCCESS";
    }

    @Override
    public String chat(ChatReq chatReq) {
        LlmService tongyiService = ModelFactory.getModel(ModelEnum.TONG_YI.getCode());
        double[] vector = tongyiService.sentence(chatReq.getQuestion());
        String collectionName = vectorService.getCollectionName();
        String vectorData = vectorService.retrieval(collectionName, vector);
        if (StringUtils.isBlank(vectorData)) {
            return "No Answer";
        }
        String prompt = LlmUtil.buildPrompt(chatReq.getQuestion(), vectorData);
        String chat = tongyiService.chat(prompt);
        log.info("question: {},  answer: {}", chatReq.getQuestion(), chat);
        return chat;
    }

    public List<ChunkResult> chunk(String fileName, String txtContent) {
        // 暂时按256固定字符数量分割
        String[] lines = StrUtil.split(txtContent, 256);
        List<ChunkResult> result = new ArrayList<>();
        AtomicInteger atomicInteger = new AtomicInteger(0);
        for (String line : lines) {
            ChunkResult chunkResult = new ChunkResult();
            chunkResult.setDocId(fileName);
            chunkResult.setContent(line);
            chunkResult.setChunkId(atomicInteger.incrementAndGet());
            result.add(chunkResult);
        }
        return result;
    }
}
